package org.third.security.cert;

import java.io.ByteArrayOutputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.KeyStore.Entry.Attribute;
import java.security.KeyStore.PasswordProtection;
import java.security.KeyStore.PrivateKeyEntry;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.Signature;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Set;
import java.util.logging.Level;
import java.util.logging.Logger;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKeyFactory;

import org.bouncycastle.crypto.digests.RIPEMD160Digest;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.Resource;
import org.third.common.utils.CodecUtils;
import org.third.common.utils.StreamUtils;

//https://github.com/OpenConext/Mujina
//https://www.programcreek.com/java-api-examples/?code=StanIsAdmin/CrashCoin/CrashCoin-master/src/common/src/be/ac/ulb/crashcoin/common/utils/Cryptography.java#
public class KeyTest {
	static DefaultResourceLoader resourceLoader = new DefaultResourceLoader();
	private static KeyStore keyStore;
	private static MessageDigest hasher = null;

	private static SecureRandom randomGenerator = null;

	/**
	 * Key derivation.
	 */
	private static RIPEMD160Digest ripemdHasher = new RIPEMD160Digest();

	private static Cipher cipher = null;

	private static SecretKeyFactory secretKeyFactory = null;

	/**
	 * DSA key pair generator.
	 */
	private static KeyPairGenerator dsaKeyGenerator = null;

	/**
	 * DSA public/private key constructor from bytes.
	 */
	private static KeyFactory dsaKeyFactory = null;
	private static Signature dsaSigner;
	static {
		InputStream inputStream = null;
		try {
			inputStream = X509CertificateTest.class.getResourceAsStream("server.p12");
			keyStore = KeyStore.getInstance("PKCS12");
			char[] password = "changeit".toCharArray();
			keyStore.load(inputStream, password);

			// Hashing
			try {
				hasher = MessageDigest.getInstance(Parameters.HASH_ALGORITHM);
			} catch (NoSuchAlgorithmException ex) {
				logAndAbort("Unable to use SHA-256 hash... Abort!", ex);
			}
			// Key derivation
			try {
				cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
			} catch (NoSuchAlgorithmException | NoSuchPaddingException ex) {
				logAndAbort("Unable to get cipher: \"AES/CBC/PKCS5Padding\". Abort!", ex);
			}
			try {
				secretKeyFactory = SecretKeyFactory.getInstance("PBKDF2WithHmacSHA1");
			} catch (NoSuchAlgorithmException ex) {
				logAndAbort("Unable to create SecretKeyFactory. Abort!", ex);
			}
			// Secure random generation
			try {
				randomGenerator = SecureRandom.getInstance("SHA1PRNG", "SUN");
			} catch (NoSuchAlgorithmException | NoSuchProviderException ex) {
				logAndAbort("Unable to create Secure Random \"SHA1PRNG\". Abort!", ex);
			}
			// DSA key generation
			try {
				dsaKeyGenerator = KeyPairGenerator.getInstance("DSA", "SUN");
			} catch (NoSuchAlgorithmException | NoSuchProviderException ex) {
				logAndAbort("unable to create DSA Keygen. Abort!", ex);
			}
			try {
				dsaKeyFactory = KeyFactory.getInstance("DSA");
			} catch (NoSuchAlgorithmException ex) {
				logAndAbort("Unable to create DSA key factory. Abort!", ex);
			}
			try {
				dsaSigner = Signature.getInstance("SHA1withDSA", "SUN");
			} catch (NoSuchAlgorithmException e) {
				logAndAbort("[Error] Could not find DSA signature algorithm. Abort!", e);
			} catch (NoSuchProviderException e) {
				logAndAbort("[Error] Could not find provider for DSA. Abort!", e);
			}
		} catch (Exception e) {
		} finally {
			StreamUtils.closeQuietly(inputStream);
		}
	}

	private static void logAndAbort(final String message, final Throwable exception) {
		Logger.getLogger(Cryptography.class.getName()).log(Level.SEVERE, message, exception);
		System.exit(1);
	}

	public static void main(String[] args) throws Exception {
		new KeyTest().startTest();
	}

	public static byte[] decrypt(byte[] encryptedBytes, PrivateKey privateKey, int keyLength, int reserveSize,
			String cipherAlgorithm) throws Exception {
		int keyByteSize = keyLength / 8;
		int decryptBlockSize = keyByteSize - reserveSize;
		int nBlock = encryptedBytes.length / keyByteSize;
		ByteArrayOutputStream outbuf = null;
		try {
			Cipher cipher = Cipher.getInstance(cipherAlgorithm);
			cipher.init(Cipher.DECRYPT_MODE, privateKey);

			outbuf = new ByteArrayOutputStream(nBlock * decryptBlockSize);
			for (int offset = 0; offset < encryptedBytes.length; offset += keyByteSize) {
				int inputLen = encryptedBytes.length - offset;
				if (inputLen > keyByteSize) {
					inputLen = keyByteSize;
				}
				byte[] decryptedBlock = cipher.doFinal(encryptedBytes, offset, inputLen);
				outbuf.write(decryptedBlock);
			}
			outbuf.flush();
			return outbuf.toByteArray();
		} catch (Exception e) {
			throw new Exception("DEENCRYPT ERROR:", e);
		} finally {
			try {
				if (outbuf != null) {
					outbuf.close();
				}
			} catch (Exception e) {
				outbuf = null;
				throw new Exception("CLOSE ByteArrayOutputStream ERROR:", e);
			}
		}
	}

	public static byte[] encrypt(byte[] plainBytes, PublicKey publicKey, int keyLength, int reserveSize,
			String cipherAlgorithm) throws Exception {
		int keyByteSize = keyLength / 8;
		int encryptBlockSize = keyByteSize - reserveSize;
		int nBlock = plainBytes.length / encryptBlockSize;
		if ((plainBytes.length % encryptBlockSize) != 0) {
			nBlock += 1;
		}
		ByteArrayOutputStream outbuf = null;
		try {
			Cipher cipher = Cipher.getInstance(cipherAlgorithm);
			cipher.init(Cipher.ENCRYPT_MODE, publicKey);

			outbuf = new ByteArrayOutputStream(nBlock * keyByteSize);
			for (int offset = 0; offset < plainBytes.length; offset += encryptBlockSize) {
				int inputLen = plainBytes.length - offset;
				if (inputLen > encryptBlockSize) {
					inputLen = encryptBlockSize;
				}
				byte[] encryptedBlock = cipher.doFinal(plainBytes, offset, inputLen);
				outbuf.write(encryptedBlock);
			}
			outbuf.flush();
			return outbuf.toByteArray();
		} catch (Exception e) {
			throw new Exception("ENCRYPT ERROR:", e);
		} finally {
			try {
				if (outbuf != null) {
					outbuf.close();
				}
			} catch (Exception e) {
				outbuf = null;
				throw new Exception("CLOSE ByteArrayOutputStream ERROR:", e);
			}
		}
	}

	private void startTest() throws NoSuchAlgorithmException, InvalidKeySpecException, IOException, Exception {
		for (java.util.Enumeration<String> aliasEnum = keyStore.aliases(); aliasEnum.hasMoreElements();) {
			String alias = aliasEnum.nextElement();
			java.security.KeyStore.ProtectionParameter protectionParam = new PasswordProtection(
					"changeit".toCharArray());
			java.security.KeyStore.Entry entry = keyStore.getEntry(alias, protectionParam);
			if (entry instanceof java.security.KeyStore.PrivateKeyEntry) {
				PrivateKeyEntry pkEntry = (PrivateKeyEntry) entry;
				// java.security.PKCS12Attribute
				Set<Attribute> pkcs12Attributes = pkEntry.getAttributes();
				Certificate endCertificate = pkEntry.getCertificate();
				java.security.PrivateKey privateKey = pkEntry.getPrivateKey();
//				System.out.println(CodecUtils.encodeBase64String(endCertificate.getEncoded()));
//				System.out.println(CodecUtils.encodeBase64String(privateKey.getEncoded()));
			}
//			System.out.println(entry);

		}
//		loadCert("/ssl/server.crt");
		loadPublicKey("/ssl/server_x509_public.key");
//		loadPrivateKey("/ssl/server_pkcs8.key");
//		generateKeyFile("test.crt", "test.key");
	}

	private void generateKeyFile(String crtFile, String keyFile)
			throws NoSuchAlgorithmException, FileNotFoundException, IOException {
		java.security.KeyPair keyPair = generateKeyPair();
		java.security.PublicKey publicKey = keyPair.getPublic();
		java.security.PrivateKey privateKey = keyPair.getPrivate();

		java.security.spec.X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(publicKey.getEncoded());
		java.security.spec.PKCS8EncodedKeySpec pkcs8Spec = new PKCS8EncodedKeySpec(privateKey.getEncoded());

		StreamUtils.writeStringToFile(CodecUtils.encodeBase64String(pkcs8Spec.getEncoded()), keyFile);
		StreamUtils.writeStringToFile(CodecUtils.encodeBase64String(x509KeySpec.getEncoded()), crtFile);
	}

	public KeyPair generateKeyPair() throws NoSuchAlgorithmException {
		KeyPairGenerator keyGen = java.security.KeyPairGenerator.getInstance("RSA");
		SecureRandom random = new SecureRandom();
		keyGen.initialize(4096, random);
		KeyPair keyPair = keyGen.generateKeyPair();
		return keyPair;
	}

	public X509Certificate loadCert(String certLocation) throws CertificateException, IOException {
		CertificateFactory cf = CertificateFactory.getInstance("X509");
		Resource certRes = resourceLoader.getResource(certLocation);
		// inputStream: DER or base64-encoded bounded at the beginning by -----BEGIN
		// CERTIFICATE-----, and must be bounded at the end by -----END
		// CERTIFICATE-----.
		X509Certificate cert = (X509Certificate) cf.generateCertificate(certRes.getInputStream());
		System.out.println(CodecUtils.encodeBase64String(cert.getEncoded()));
		System.out.println(CodecUtils.encodeBase64String((cert.getPublicKey().getEncoded())));
		return cert;
	}

	/**
	 * Given a resource location it loads a DER RSA private Key.
	 *
	 * @param privateKeyLocation
	 * @return
	 * @throws IOException
	 * @throws NoSuchAlgorithmException
	 * @throws InvalidKeySpecException
	 */
	public RSAPrivateKey loadPrivateKey(String privateKeyLocation)
			throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {
		Resource keyRes = resourceLoader.getResource(privateKeyLocation);
		String keyString = StreamUtils.pkiStreamToString(keyRes.getInputStream());
		System.out.println(keyString);
		byte[] keyBytes = CodecUtils.decodeBase64(keyString);
		PKCS8EncodedKeySpec privateKeySpec = new PKCS8EncodedKeySpec(keyBytes);
		KeyFactory keyFactory = KeyFactory.getInstance("RSA");
		RSAPrivateKey privateKey = (RSAPrivateKey) keyFactory.generatePrivate(privateKeySpec);
		return privateKey;
	}

	public RSAPublicKey loadPublicKey(String publicKeyLocation)
			throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {
		Resource keyRes = resourceLoader.getResource(publicKeyLocation);
		String keyString = StreamUtils.pkiStreamToString(keyRes.getInputStream());
		System.out.println(keyString);
		byte[] keyBytes = CodecUtils.decodeBase64(keyString);
		X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(keyBytes);
		KeyFactory keyFactory = KeyFactory.getInstance("RSA");
		RSAPublicKey publicKey = (RSAPublicKey) keyFactory.generatePublic(publicKeySpec);
		return publicKey;
	}
}
